#include <stdio.h>
#include <cuda_runtime.h>

// vector add
__global__ void add(int *a, int *b, int *c, int n) {
	int index = threadIdx.x + blockIdx.x * blockDim.x;
	if (index < n) {
		c[index] = a[index] + b[index];
	}
}

int main() {
	int n = 256; // size of vectors
	size_t bytes = n * sizeof(int);
	int *h_a, *h_b, *h_c; // host vectors
	int *d_a, *d_b, *d_c; // device vectors

	// allocate host memory
	h_a = (int *)malloc(bytes);
	h_b = (int *)malloc(bytes);
	h_c = (int *)malloc(bytes);

	// init host vectors
	for (int i = 0; i < n; i++) {
		h_a[i] = i;
		h_b[i] = i;
	}

	// allocate device memory
	cudaMalloc(&d_a, bytes);
	cudaMalloc(&d_b, bytes);
	cudaMalloc(&d_c, bytes);

	// create 2 streams
	cudaStream_t stream1, stream2;
	cudaStreamCreateWithFlags(&stream1, cudaStreamNonBlocking);
	cudaStreamCreateWithFlags(&stream2, cudaStreamNonBlocking);

	// asynchronously copy memory to the device on stream1
	cudaMemcpyAsync(d_a, h_a, bytes, cudaMemcpyHostToDevice, stream1);
	cudaMemcpyAsync(d_b, h_b, bytes, cudaMemcpyHostToDevice, stream1);

	cudaStreamSynchronize(stream1);

	// execute the kernel on stream2
	int threads = 256;
	int blocks = (n + threads - 1) / threads;
	add<<<blocks, threads, 0, stream2>>>(d_a, d_b, d_c, n);

	// wait for the memory copy on stream1 to complete
	cudaStreamSynchronize(stream1);

	// asynchronously copy the result back to the host on stream2
	cudaMemcpyAsync(h_c, d_c, bytes, cudaMemcpyDeviceToHost, stream2);

	// wait for the memory copy on stream2 to complete
	cudaStreamSynchronize(stream2);

	// check the results
	int i;
	for (i = 0; i < n; i++) {
		if (h_c[i] != 2 * i) {
			printf("Error: invalid result at %d\n", i);
			break;
		}
	}

	if (i >= n) {
		printf("The result is correct.\n");
	}

	// clean up resources
	cudaStreamDestroy(stream1);
	cudaStreamDestroy(stream2);
	cudaFree(d_a);
	cudaFree(d_b);
	cudaFree(d_c);
	free(h_a);
	free(h_b);
	free(h_c);

	return 0;
}
